import redis
import json
import time
import numpy as np
from datetime import datetime
from sklearn.metrics.pairwise import cosine_similarity
from typing import List, Dict, Tuple, Optional


class RedisChatManager:
    """使用Redis管理大模型多轮对话记录的类"""

    def __init__(self, host='localhost', port=6379, db=0,
                 decay_factor=0.001, topk_coarse=2000):
        """
        初始化Redis连接和参数设置

        :param host: Redis主机地址
        :param port: Redis端口
        :param db: Redis数据库编号
        :param decay_factor: 时间衰减因子
        :param topk_coarse: 粗筛保留的对话数量
        """
        self.redis_client = redis.Redis(host=host, port=port, db=db, decode_responses=True)
        self.decay_factor = decay_factor  # 时间衰减因子
        self.topk_coarse = topk_coarse  # 粗筛保留的对话数量

        # 用于存储对话的Redis键前缀
        self.chat_prefix = "chat:"
        self.timestamp_set = "chat:timestamps"  # 存储所有对话的时间戳有序集合

    def add_dialogue(self, session_id: str, user_msg: str, bot_msg: str,
                     user_embedding: Optional[List[float]] = None,
                     bot_embedding: Optional[List[float]] = None) -> str:
        """
        添加新的对话记录

        :param session_id: 会话ID
        :param user_msg: 用户消息
        :param bot_msg: 机器人回复
        :param user_embedding: 用户消息的嵌入向量
        :param bot_embedding: 机器人消息的嵌入向量
        :return: 对话记录ID
        """
        timestamp = time.time()
        dialogue_id = f"{session_id}:{int(timestamp * 1000)}"
        key = f"{self.chat_prefix}{dialogue_id}"

        # 存储对话内容
        dialogue_data = {
            "session_id": session_id,
            "user_msg": user_msg,
            "bot_msg": bot_msg,
            "timestamp": timestamp,
            "user_embedding": json.dumps(user_embedding) if user_embedding else None,
            "bot_embedding": json.dumps(bot_embedding) if bot_embedding else None
        }

        # 过滤掉None值
        dialogue_data = {k: v for k, v in dialogue_data.items() if v is not None}

        # 存储到Redis
        self.redis_client.hset(key, mapping=dialogue_data)
        # 添加时间戳到有序集合
        self.redis_client.zadd(self.timestamp_set, {dialogue_id: timestamp})

        # 保持集合大小不超过粗筛阈值的1.5倍，防止集合过大
        total_count = self.redis_client.zcard(self.timestamp_set)
        if total_count > self.topk_coarse * 1.5:
            # 移除最旧的记录
            self.redis_client.zremrangebyrank(self.timestamp_set, 0,
                                              total_count - self.topk_coarse * 1.5)

        return dialogue_id

    def _get_time_decay_weight(self, timestamp: float) -> float:
        """
        计算时间衰减权重

        :param timestamp: 对话时间戳
        :return: 衰减后的权重
        """
        current_time = time.time()
        time_diff = current_time - timestamp
        # 指数衰减: weight = e^(-decay_factor * time_diff)
        return np.exp(-self.decay_factor * time_diff)

    def _get_semantic_similarity(self, embedding1: List[float], embedding2: List[float]) -> float:
        """
        计算两个嵌入向量的语义相似度

        :param embedding1: 第一个嵌入向量
        :param embedding2: 第二个嵌入向量
        :return: 相似度分数(0-1)
        """
        if not embedding1 or not embedding2:
            return 0.0

        # 转换为二维数组以适应cosine_similarity函数
        vec1 = np.array(embedding1).reshape(1, -1)
        vec2 = np.array(embedding2).reshape(1, -1)
        return cosine_similarity(vec1, vec2)[0][0]

    def get_topk_dialogues(self, session_id: Optional[str] = None,
                           query_embedding: Optional[List[float]] = None,
                           topk: int = 10, time_weight: float = 0.5,
                           semantic_weight: float = 0.5) -> List[Dict]:
        """
        获取TopK的对话记录，基于时间衰减和语义相似度的加权

        :param session_id: 可选，指定会话ID筛选
        :param query_embedding: 查询的嵌入向量，用于语义相似度计算
        :param topk: 要返回的对话数量
        :param time_weight: 时间权重因子
        :param semantic_weight: 语义权重因子
        :return: 排序后的对话记录列表
        """
        # 1. 粗筛：获取最近的topk_coarse条对话
        recent_dialogue_ids = self.redis_client.zrevrange(
            self.timestamp_set, 0, self.topk_coarse - 1)

        # 2. 加载对话详情并计算权重
        dialogue_list = []
        for dialogue_id in recent_dialogue_ids:
            key = f"{self.chat_prefix}{dialogue_id}"
            dialogue_data = self.redis_client.hgetall(key)

            # 如果指定了会话ID，只保留对应会话的对话
            if session_id and dialogue_data.get("session_id") != session_id:
                continue

            # 解析数据
            try:
                timestamp = float(dialogue_data["timestamp"])
                user_embedding = json.loads(
                    dialogue_data["user_embedding"]) if "user_embedding" in dialogue_data else None
                bot_embedding = json.loads(dialogue_data["bot_embedding"]) if "bot_embedding" in dialogue_data else None

                # 计算时间衰减权重
                time_weight_val = self._get_time_decay_weight(timestamp)

                # 计算语义相似度权重
                semantic_sim = 0.0
                if query_embedding:
                    # 取用户消息和机器人消息中与查询更相似的那个
                    user_sim = self._get_semantic_similarity(user_embedding, query_embedding) if user_embedding else 0.0
                    bot_sim = self._get_semantic_similarity(bot_embedding, query_embedding) if bot_embedding else 0.0
                    semantic_sim = max(user_sim, bot_sim)

                # 计算综合权重
                combined_weight = (time_weight * time_weight_val) + (semantic_weight * semantic_sim)

                # 添加到列表
                dialogue_list.append({
                    "dialogue_id": dialogue_id,
                    "session_id": dialogue_data["session_id"],
                    "user_msg": dialogue_data["user_msg"],
                    "bot_msg": dialogue_data["bot_msg"],
                    "timestamp": timestamp,
                    "time_weight": time_weight_val,
                    "semantic_similarity": semantic_sim,
                    "combined_weight": combined_weight,
                    "datetime": datetime.fromtimestamp(timestamp).strftime('%Y-%m-%d %H:%M:%S')
                })
            except Exception as e:
                print(f"处理对话 {dialogue_id} 时出错: {e}")
                continue

        # 3. 按综合权重排序并返回TopK
        dialogue_list.sort(key=lambda x: x["combined_weight"], reverse=True)
        return dialogue_list[:topk]

    def clear_session(self, session_id: str) -> int:
        """
        清除指定会话的所有对话记录

        :param session_id: 会话ID
        :return: 清除的记录数量
        """
        # 这是一个简单实现，实际应用中可能需要优化
        count = 0
        all_dialogue_ids = self.redis_client.zrange(self.timestamp_set, 0, -1)

        for dialogue_id in all_dialogue_ids:
            if dialogue_id.startswith(f"{session_id}:"):
                key = f"{self.chat_prefix}{dialogue_id}"
                self.redis_client.delete(key)
                self.redis_client.zrem(self.timestamp_set, dialogue_id)
                count += 1

        return count

    def get_session_dialogues(self, session_id: str) -> List[Dict]:
        """获取指定会话的所有对话记录（按时间排序）"""
        dialogues = []
        all_dialogue_ids = self.redis_client.zrange(self.timestamp_set, 0, -1)

        for dialogue_id in all_dialogue_ids:
            if dialogue_id.startswith(f"{session_id}:"):
                key = f"{self.chat_prefix}{dialogue_id}"
                dialogue_data = self.redis_client.hgetall(key)
                if dialogue_data:
                    try:
                        timestamp = float(dialogue_data["timestamp"])
                        dialogues.append({
                            "dialogue_id": dialogue_id,
                            "user_msg": dialogue_data["user_msg"],
                            "bot_msg": dialogue_data["bot_msg"],
                            "timestamp": timestamp,
                            "datetime": datetime.fromtimestamp(timestamp).strftime('%Y-%m-%d %H:%M:%S')
                        })
                    except Exception as e:
                        print(f"处理对话 {dialogue_id} 时出错: {e}")

        # 按时间排序
        dialogues.sort(key=lambda x: x["timestamp"])
        return dialogues


# 使用示例
if __name__ == "__main__":
    # 初始化对话管理器
    chat_manager = RedisChatManager(decay_factor=0.0001, topk_coarse=2000)

    # 示例：添加一些对话记录
    session_id = "user_123"


    # 生成随机嵌入向量（实际应用中应使用模型生成）
    def generate_dummy_embedding():
        return list(np.random.rand(384))  # 假设384维的嵌入向量


    # 添加几条对话
    for i in range(15):
        user_msg = f"这是第{i + 1}个问题：关于人工智能的应用"
        bot_msg = f"这是第{i + 1}个回答：人工智能在多个领域有广泛应用"

        # 模拟时间间隔
        if i > 0:
            time.sleep(0.1)

        chat_manager.add_dialogue(
            session_id=session_id,
            user_msg=user_msg,
            bot_msg=bot_msg,
            user_embedding=generate_dummy_embedding(),
            bot_embedding=generate_dummy_embedding()
        )

    # 示例查询：获取与特定查询最相关的Top5对话
    query_embedding = generate_dummy_embedding()  # 实际应用中应使用查询文本的嵌入
    top_dialogues = chat_manager.get_topk_dialogues(
        session_id=session_id,
        query_embedding=query_embedding,
        topk=5,
        time_weight=0.3,
        semantic_weight=0.7
    )

    print("Top 5 对话记录:")
    for i, dialogue in enumerate(top_dialogues, 1):
        print(f"\n第{i}条 (权重: {dialogue['combined_weight']:.4f}):")
        print(f"时间: {dialogue['datetime']}")
        print(f"用户: {dialogue['user_msg']}")
        print(f"机器人: {dialogue['bot_msg']}")
        print(f"时间权重: {dialogue['time_weight']:.4f}, 语义相似度: {dialogue['semantic_similarity']:.4f}")

    # 获取整个会话的对话历史
    session_history = chat_manager.get_session_dialogues(session_id)
    print(f"\n会话 {session_id} 共有 {len(session_history)} 条对话记录")

    # 清除会话（实际应用中谨慎使用）
    # cleared = chat_manager.clear_session(session_id)
    # print(f"已清除 {cleared} 条对话记录")